# -*- coding: utf-8 -*-
"""
Created on Thu Apr 21 23:09:03 2022

@author: 13503
"""

import os
import torch
import random
import numpy as np
import h5py

"""
@brief: The mat scheme file is in dataset and transform them into dataset_cache and store in dataset_cache directory 
"""
class SignNetData():
    def __init__(self, file_path:str, train_ratio:int=0.7, beta:float=0.1, gamma:float=10):
        """
        parameters:
            @file_path: define the relative path of data to load
            @train_ratio: decide the percentage of data used for training
        """
        print("using {0}% data to train model".format(train_ratio*100))
        save_dir = os.path.split(file_path)[0]
        self.file_path = file_path
        self.train_ratio = train_ratio
        self.full_mat_save_path = os.path.join(save_dir, "full_" + os.path.split(file_path)[-1].split(sep=".")[0] + ".pt")
        self.train_mat_save_path = os.path.join(save_dir, "train_" + os.path.split(file_path)[-1].split(sep=".")[0] + ".pt")
        self.validate_mat_save_path = os.path.join(save_dir, "validate_" + os.path.split(file_path)[-1].split(sep=".")[0] + ".pt")
        self.test_mat_save_path = os.path.join(save_dir, "test_" + os.path.split(file_path)[-1].split(sep=".")[0] + ".pt")
        self.penalty_mat_save_path = os.path.join(save_dir, "penalty_" + os.path.split(file_path)[-1].split(sep=".")[0] + ".pt")

        if os.path.exists(self.full_mat_save_path):
            full_data = torch.load(self.full_mat_save_path)
        else:
            full_data = self.load_tensor_from_mat(self.file_path)
            torch.save(full_data, self.full_mat_save_path)

       
        if not (os.path.exists(self.train_mat_save_path) 
           and  os.path.exists(self.validate_mat_save_path) 
           and  os.path.exists(self.penalty_mat_save_path)):

            print("begin to split dataset")
            nz_idx = torch.nonzero(full_data)
            random.shuffle(nz_idx)
            num_of_train_sample = int(self.train_ratio * len(nz_idx))
            num_of_val_sample = int((len(nz_idx) - num_of_train_sample) / 2)

            train_data = torch.zeros(full_data.shape)
            validate_data = torch.zeros(full_data.shape)
            test_data = torch.zeros(full_data.shape)
            penalty_mat = torch.ones(full_data.shape)
            print("There are {0} edges in the graph.".format(len(nz_idx)))
            
            idx_cnt = 0
            for coord in nz_idx:
                coord = tuple(coord)
                if idx_cnt % 100000 == 0:
                    print("already processed {0}/{1} edges".format(idx_cnt, len(nz_idx)))
                if idx_cnt < num_of_train_sample:
                    train_data[coord] = full_data[coord]    
                elif num_of_train_sample <= idx_cnt < num_of_train_sample + num_of_val_sample:
                    validate_data[coord] = full_data[coord]
                else:
                    test_data[coord] = full_data[coord]
                
                if full_data[coord] == 1:
                    penalty_mat[coord] = beta
                else:
                    penalty_mat[coord] = beta * gamma
                idx_cnt = idx_cnt + 1

            if not os.path.exists(save_dir):
                os.mkdir(save_dir)
            
            torch.save(train_data, self.train_mat_save_path)
            torch.save(validate_data, self.validate_mat_save_path)
            torch.save(test_data, self.test_mat_save_path)
            print("self.penalty_mat_save_path = {0}".format(self.penalty_mat_save_path))
            torch.save(penalty_mat, self.penalty_mat_save_path)
            print("finish split dataset")


    """
    @brief: This is a function to read data from .mat file into tensor
    """
    def load_tensor_from_mat(self, load_path:str)->torch.tensor:
        if os.path.exists(load_path):
            np_data = np.transpose(np.array(h5py.File(name=load_path)["Gwl_ud"])).astype(np.float16)
            data = torch.from_numpy(np_data).type(torch.float) # transform array into tensor
            return data
        else:
            raise OSError("The input argument file_location does not exist!")

    def get_full_data(self):
        if os.path.exists(self.full_mat_save_path):
            raise OSError("The full_mat_save_path does not exist!")
        return self.load_tensor_from_mat(self.full_mat_save_path)

    def get_train_data(self):
        if os.path.exists(self.train_mat_save_path):
            raise OSError("The train_mat_save_path does not exist!")
        return self.load_tensor_from_mat(self.train_mat_save_path)
        
    def get_validate_data(self):
        if os.path.exists(self.validate_mat_save_path):
            raise OSError("The validate_mat_save_path does not exist!")
        return self.load_tensor_from_mat(self.validate_mat_save_path)

    def get_test_data(self):
        if os.path.exists(self.test_mat_save_path):
            raise OSError("The test_mat_save_path does not exist!")
        return self.load_tensor_from_mat(self.test_mat_save_path)